# coding=utf-8
# @author:      ChengJing
# @name:        IFA.py
# @datetime:    2022/2/3 0:23
# @software:    PyCharm
# @description:

import torch
import torch.nn as nn
from model.improve_location.ImproveFeatures import LAQ


class SELayer(nn.Module):
    """
    注意力模块
    """
    def __init__(self, features, reduction=16):
        """
        Args:
            features: 输入特征数
            reduction: 缩小倍数
        """
        super(SELayer, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(features, features // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(features // reduction, features, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        """
        前向传递函数
         """
        y = self.fc(x)
        return x * y.expand_as(x)


class Block(nn.Module):
    """
    FADenseNet块的定义
    """
    def __init__(self, in_features, hidden1, hidden2, hidden3, out_features, dropout=0.5):
        """
        Args:
            in_features:int, 输入特征数
            hidden1:int, 隐含层特征数
            hidden2:int, 隐含层特征数
            hidden3:int, 隐含层特征数
            out_features:int, 输出特征数
            dropout: float, dropout的概率值
        """
        super(Block, self).__init__()
        self.short = nn.Sequential(
            nn.BatchNorm1d(in_features),
            nn.ReLU()
        )
        self.density = nn.Sequential(
            nn.Linear(in_features, hidden1),
            nn.ReLU(),
            nn.Linear(hidden1, hidden2),
            nn.ReLU(),
            nn.Linear(hidden2, hidden3),
            nn.ReLU(),
            nn.Linear(hidden3, out_features),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        """
        前向传递函数
        """
        x = self.density(self.short(x))
        return x


class IFA(nn.Module):
    """
    FADenseNet模型
    """
    def __init__(self, g, sensors, in_features, out_features):
        """
        Args:
            g: 流量监测点和压力监测点之间的连接关系，shape：N*M
            sensors: 压力监测点的数量
            in_features: int,输入特征数
            out_features: int,输出特征数
        """
        super(IFA, self).__init__()
        self.laq = LAQ(g, sensors, 2*sensors)
        self.input = nn.Linear(in_features, 49)
        self.block1 = Block(49, 120, 206, 124, 49)
        self.se1 = SELayer(49, 8)
        self.block2 = Block(49, 144, 248, 144, 49)
        self.se2 = SELayer(49, 8)
        self.block3 = Block(49, 155, 312, 155, 49)
        self.se3 = SELayer(49, 8)
        self.block4 = Block(49, 186, 366, 186, 49)
        self.se4 = SELayer(49, 8)
        self.block5 = Block(49, 220, 111, 68, out_features)
        self.output = nn.Softmax(dim=1)

    def forward(self, data):
        """
        前向传递函数
        """
        x, q = data
        x = self.laq(x, q)
        x = x.reshape((x.shape[0], -1))
        x = self.input(x)
        d1 = x
        x = self.block1(d1)
        x = self.se1(x)
        d2 = x + d1
        x = self.block2(d2)
        x = self.se2(x)
        d3 = x + d1 + d2
        x = self.block3(d3)
        x = self.se3(x)
        d4 = x + d1 + d2 + d3
        x = self.block4(d4)
        x = self.se4(x)
        d5 = x + d1 + d2 + d3 + d4
        x = self.block5(d5)
        x = self.output(x)
        return x


if __name__ == '__main__':
    x = torch.rand((10, 20, 6))
    q = torch.rand((10, 3))
    g = torch.rand((3, 6))
    model = IFA(g, 6, 6*20, 20)
    y = model((x, q))
    print(y.shape)
